#ifndef EST_POLYA_MOMENT_HEADER
#define EST_POLYA_MOMENT_HEADER

#include<string>
#include<fstream>
#include<sstream>
#include<iostream>
#include<algorithm>
#include<iterator>
#include<vector>
#include<map>
#include<cstdlib>
#include<ctime>
#include<cmath>
#include <gsl/gsl_rng.h>
#include <gsl/gsl_randist.h>
#include <gsl/gsl_sort.h>
#include <gsl/gsl_statistics.h>
#include "def.hpp"
#include "utils.hpp"
#include "io.hpp"

using namespace std;

void log_evidence(double**& a, double*& estimated_p, int& K, int& N){
  for(int k = 0; k < K; k++){
    estimated_p[k] = 0.0;
  }

  for(int i=0; i < N; i++){
    for(int k =0; k < K; k++){
      estimated_p[k] = estimated_p[k] + log(a[i][k]);
    }
  }
  for(int k =0; k < K; k++){
    estimated_p[k] = estimated_p[k] / N;
  }
}

void moment_matching(int**& a, int*& b, double*& estimated_alpha, int& K, int& N){
  double* mean = new double[K];
  double* second_moment = new double[K];
  double* s = new double[K];
  // clean array
  for(int k = 0; k < K; k ++){
    mean[k] = 0.0;
    second_moment[k] = 0.0;
    s[k] = 0.0;
  }  
  
  for(int i = 0; i < N ; i ++ ){
    for(int k = 0; k < K ; k ++ ){
      double temp = (a[i][k] + 1.0)/(b[i] + K);
      mean[k] = mean[k] + temp;
      second_moment[k] = second_moment[k] + temp * temp;
    }
  }

  for(int k = 0; k < K; k ++){
    mean[k] = mean[k] / N;
    second_moment[k] = second_moment[k] / N;
  }

  for(int k = 0; k < K; k++){
    s[k] = (mean[k] - second_moment[k])/ (second_moment[k] - mean[k] * mean[k]);
  }

  // get the median
  gsl_sort (s, 1, K);
  double estimated_s = gsl_stats_median_from_sorted_data(s, 1, K);

  for(int k = 0; k < K; k++){
    estimated_alpha[k] = mean[k] * estimated_s;
  }

  delete[] mean;
  delete[] second_moment;
  delete[] s;
}

#endif
